import torch
import os
import numpy as np
from matplotlib import pyplot as plt
from torchvision import datasets,transforms
from torch.utils.data import DataLoader

def data_loader(root='./data', batch_size = 32, spilt = 0.7):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ]);
    train_dataset = datasets.CIFAR10(root=root, train=True, download=True, transform=transform)
    test_dataset = datasets.CIFAR10(root=root, train=False, download=True, transform=transform)
    datasets_size = len(train_dataset)
    train_size, val_size = int(datasets_size * spilt), int(datasets_size * (1-spilt))
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
    return train_loader, val_loader, test_loader

if __name__ == '__main__':
    train_loader, val_loader, test_loader = data_loader()
    for i, (images, labels) in enumerate(train_loader):
        print(images.size(), images.requires_grad)
        print(labels.size(), labels.requires_grad)
        break


